import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
import argparse
import os
import json
import time
from sklearn.metrics import mean_squared_error, mean_absolute_error
from sklearn.preprocessing import StandardScaler
import ast
from spikingjelly.activation_based import neuron, functional, layer, encoding

# 超参数 - 按照USAGE文档设置
TIME_STEPS = 14
DELAY_STEPS = 5
LIQUID_NEURONS = 135  # 135个I&F神经元
READOUT_NEURONS = 51  # 51个读出神经元
INPUT_SIZE = 14
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# LSM模型参数
LSM_PARAMS = {
    # 神经元参数
    "membrane_time_constant": 30.0,  # ms
    "excitatory_refractory": 3.0,    # ms
    "inhibitory_refractory": 2.0,    # ms
    "threshold_potential": 15.0,     # mV
    "reset_voltage": 13.5,           # mV
    "background_current": 13.5,      # nA
    "input_resistance": 1.0,         # MΩ
    
    # 网络结构
    "inhibitory_ratio": 0.2,         # 20%抑制性神经元
    "grid_size": (15, 3, 3),         # 15×3×3空间网格
    
    # 连接概率参数
    "connection_params": {
        "EE": {"C": 0.3, "lambda": 2.0},
        "EI": {"C": 0.2, "lambda": 2.0},
        "IE": {"C": 0.4, "lambda": 2.0},
        "II": {"C": 0.1, "lambda": 2.0}
    },
    
    # 动态突触参数 (Markram模型)
    "synapse_params": {
        "EE": {"U": 0.5, "D": 1.1, "F": 0.05},
        "EI": {"U": 0.05, "D": 0.125, "F": 1.2},
        "IE": {"U": 0.25, "D": 0.7, "F": 0.02},
        "II": {"U": 0.32, "D": 0.144, "F": 0.06}
    },
    
    # 突触强度参数 (nA)
    "synapse_strength": {
        "EE": 30.0,
        "EI": 60.0,
        "IE": -19.0,
        "II": -19.0,
        "input_excitatory": 18.0,
        "input_inhibitory": 9.0
    },
    
    # 突触后电流时间常数
    "synaptic_time_constants": {
        "excitatory": 3.0,  # ms
        "inhibitory": 6.0   # ms
    },
    
    # 传输延迟
    "transmission_delays": {
        "EE": 1.5,  # ms
        "other": 0.8  # ms
    }
}

class DynamicSynapse:
    """动态突触模型 (Markram模型)"""
    
    def __init__(self, U, D, F, std_ratio=0.5):
        self.U = U + torch.randn(1) * U * std_ratio  # 利用率
        self.D = D + torch.randn(1) * D * std_ratio  # 恢复时间常数
        self.F = F + torch.randn(1) * F * std_ratio  # 促进时间常数
        self.R = 1.0  # 可用资源
        self.u = 0.0  # 利用率状态
        
    def update(self, spike):
        """更新突触状态"""
        if spike:
            # 脉冲到达
            self.u = self.u + self.U * (1 - self.u)
            release = self.u * self.R
            self.R = self.R - release
            self.u = self.u * (1 - self.U)
        else:
            # 无脉冲
            self.R = self.R + (1 - self.R) / self.D
            self.u = self.u * (1 - self.U)
            
        return release if spike else 0.0

class BioPhysicalLSM(nn.Module):
    """生物物理精确的LSM模型"""
    
    def __init__(self, input_size, liquid_neurons, readout_neurons, params):
        super().__init__()
        self.input_size = input_size
        self.liquid_neurons = liquid_neurons
        self.readout_neurons = readout_neurons
        self.params = params
        
        # 确定兴奋性和抑制性神经元
        self.inhibitory_mask = self._create_inhibitory_mask()
        
        # 创建空间位置
        self.neuron_positions = self._create_neuron_positions()
        
        # 创建连接矩阵
        self.connection_matrix = self._create_connection_matrix()
        
        # 创建突触权重
        self.synapse_weights = self._create_synapse_weights()
        
        # 创建动态突触
        self.dynamic_synapses = self._create_dynamic_synapses()
        
        # 创建传输延迟
        self.transmission_delays = self._create_transmission_delays()
        
        # 输入连接
        self.input_weights = self._create_input_weights()
        
        # 不应期状态
        self.refractory_states = torch.zeros(liquid_neurons)
        
        # 使用IF神经元，设置生物物理参数
        self.neurons = neuron.IFNode(
            v_threshold=params["threshold_potential"],
            v_reset=params["reset_voltage"]
        )
        # 膜时间常数通过其他方式实现
        self.membrane_time_constant = params["membrane_time_constant"]
        
        # 读出层
        self.readout = nn.Linear(liquid_neurons, 1)
        
    def to(self, device):
        """重写to方法，确保所有张量都移动到指定设备"""
        super().to(device)
        self.inhibitory_mask = self.inhibitory_mask.to(device)
        self.neuron_positions = self.neuron_positions.to(device)
        self.connection_matrix = self.connection_matrix.to(device)
        self.synapse_weights = self.synapse_weights.to(device)
        return self
        
    def _create_inhibitory_mask(self):
        """创建抑制性神经元掩码"""
        inhibitory_count = int(self.liquid_neurons * self.params["inhibitory_ratio"])
        mask = torch.zeros(self.liquid_neurons, dtype=torch.bool)
        inhibitory_indices = torch.randperm(self.liquid_neurons)[:inhibitory_count]
        mask[inhibitory_indices] = True
        return mask
        
    def _create_neuron_positions(self):
        """创建神经元在3D网格中的位置"""
        grid_size = self.params["grid_size"]
        positions = []
        for i in range(grid_size[0]):
            for j in range(grid_size[1]):
                for k in range(grid_size[2]):
                    positions.append([i, j, k])
        return torch.tensor(positions[:self.liquid_neurons], dtype=torch.float32)
        
    def _distance(self, pos1, pos2):
        """计算两个位置之间的欧几里得距离"""
        return torch.sqrt(torch.sum((pos1 - pos2) ** 2))
        
    def _connection_probability(self, pos1, pos2, conn_type):
        """计算连接概率"""
        distance = self._distance(pos1, pos2)
        C = self.params["connection_params"][conn_type]["C"]
        lambda_val = self.params["connection_params"][conn_type]["lambda"]
        return C * torch.exp(-(distance / lambda_val) ** 2)
        
    def _create_connection_matrix(self):
        """创建连接矩阵"""
        conn_matrix = torch.zeros(self.liquid_neurons, self.liquid_neurons)
        
        for i in range(self.liquid_neurons):
            for j in range(self.liquid_neurons):
                if i == j:
                    continue
                    
                # 确定连接类型
                i_inhibitory = self.inhibitory_mask[i]
                j_inhibitory = self.inhibitory_mask[j]
                
                if not i_inhibitory and not j_inhibitory:
                    conn_type = "EE"
                elif not i_inhibitory and j_inhibitory:
                    conn_type = "EI"
                elif i_inhibitory and not j_inhibitory:
                    conn_type = "IE"
                else:
                    conn_type = "II"
                    
                # 计算连接概率
                prob = self._connection_probability(
                    self.neuron_positions[i], 
                    self.neuron_positions[j], 
                    conn_type
                )
                
                # 根据概率决定是否连接
                if torch.rand(1) < prob:
                    conn_matrix[i, j] = 1.0
                    
        return conn_matrix
        
    def _create_synapse_weights(self):
        """创建突触权重"""
        weights = torch.zeros(self.liquid_neurons, self.liquid_neurons)
        
        for i in range(self.liquid_neurons):
            for j in range(self.liquid_neurons):
                if self.connection_matrix[i, j] == 0:
                    continue
                    
                # 确定连接类型
                i_inhibitory = self.inhibitory_mask[i]
                j_inhibitory = self.inhibitory_mask[j]
                
                if not i_inhibitory and not j_inhibitory:
                    strength = self.params["synapse_strength"]["EE"]
                elif not i_inhibitory and j_inhibitory:
                    strength = self.params["synapse_strength"]["EI"]
                elif i_inhibitory and not j_inhibitory:
                    strength = self.params["synapse_strength"]["IE"]
                else:
                    strength = self.params["synapse_strength"]["II"]
                    
                # 添加随机变异
                variation = torch.randn(1) * 0.1 * strength
                weights[i, j] = strength + variation
                
        return weights
        
    def _create_dynamic_synapses(self):
        """创建动态突触"""
        synapses = {}
        for i in range(self.liquid_neurons):
            for j in range(self.liquid_neurons):
                if self.connection_matrix[i, j] == 0:
                    continue
                    
                # 确定连接类型
                i_inhibitory = self.inhibitory_mask[i]
                j_inhibitory = self.inhibitory_mask[j]
                
                if not i_inhibitory and not j_inhibitory:
                    conn_type = "EE"
                elif not i_inhibitory and j_inhibitory:
                    conn_type = "EI"
                elif i_inhibitory and not j_inhibitory:
                    conn_type = "IE"
                else:
                    conn_type = "II"
                    
                # 获取突触参数
                params = self.params["synapse_params"][conn_type]
                synapse = DynamicSynapse(
                    U=params["U"],
                    D=params["D"], 
                    F=params["F"]
                )
                synapses[(i, j)] = synapse
                
        return synapses
        
    def _create_transmission_delays(self):
        """创建传输延迟"""
        delays = torch.zeros(self.liquid_neurons, self.liquid_neurons)
        
        for i in range(self.liquid_neurons):
            for j in range(self.liquid_neurons):
                if self.connection_matrix[i, j] == 0:
                    continue
                    
                # 确定连接类型
                i_inhibitory = self.inhibitory_mask[i]
                j_inhibitory = self.inhibitory_mask[j]
                
                if not i_inhibitory and not j_inhibitory:
                    delay = self.params["transmission_delays"]["EE"]
                else:
                    delay = self.params["transmission_delays"]["other"]
                    
                delays[i, j] = delay
                
        return delays
        
    def _create_input_weights(self):
        """创建输入权重"""
        input_weights = torch.zeros(self.input_size, self.liquid_neurons)
        
        for i in range(self.input_size):
            for j in range(self.liquid_neurons):
                if self.inhibitory_mask[j]:
                    strength = self.params["synapse_strength"]["input_inhibitory"]
                else:
                    strength = self.params["synapse_strength"]["input_excitatory"]
                    
                # 添加随机变异
                variation = torch.randn(1) * 0.1 * strength
                input_weights[i, j] = strength + variation
                
        return nn.Parameter(input_weights)
        
    def forward(self, x):
        """前向传播"""
        batch_size, T, _ = x.shape
        
        # 确保所有张量都在正确的设备上
        device = x.device
        
        # 初始化膜电位
        v = torch.rand(batch_size, self.liquid_neurons, device=device) * 1.5 + 13.5  # [13.5, 15.0] mV
        
        states = []
        
        for t in range(T):
            # 输入电流
            input_current = torch.matmul(x[:, t], self.input_weights)
            
            # 递归电流
            if t > 0:
                recurrent_current = torch.matmul(states[-1], self.synapse_weights.to(device).T)
            else:
                recurrent_current = torch.zeros(batch_size, self.liquid_neurons, device=device)
                
            # 总电流
            total_current = input_current + recurrent_current + self.params["background_current"]
            
            # 更新膜电位
            dv_dt = (self.params["reset_voltage"] - v) / self.params["membrane_time_constant"] + total_current
            v = v + dv_dt
            
            # 发放脉冲
            spikes = (v >= self.params["threshold_potential"]).float()
            
            # 重置发放神经元的膜电位
            v = v * (1 - spikes) + spikes * self.params["reset_voltage"]
            
            states.append(spikes)
            
        # 返回所有时间步的状态
        return torch.stack(states, dim=1)

class VoltageEncoder(nn.Module):
    """电压编码器：将电压值编码为脉冲序列"""
    
    def __init__(self, T):
        super().__init__()
        self.T = T

    def forward(self, x):
        # x: [batch, 14]
        # 输出: [batch, T, 14]
        x = x.unsqueeze(1).repeat(1, self.T, 1)  # 复制到每个时间步
        rate = torch.sigmoid(x)  # 将电压转换为发射率
        spikes = torch.rand_like(rate) < rate  # 生成脉冲
        return spikes.float()

class BatteryLSM(nn.Module):
    """完整的电池容量预测LSM模型"""
    
    def __init__(self, input_size, liquid_neurons, readout_neurons, T, delay_steps, params):
        super().__init__()
        self.T = T
        self.delay_steps = delay_steps
        self.params = params
        
        self.encoder = VoltageEncoder(T)
        self.liquid = BioPhysicalLSM(input_size, liquid_neurons, readout_neurons, params)
        self.readout = nn.Linear(liquid_neurons, 1)
        
    def forward(self, x):
        # 编码
        spikes = self.encoder(x)  # [batch, T, input_size]
        
        # 液体层处理
        liquid_states = self.liquid(spikes)  # [batch, T, liquid_neurons]
        
        # 在输入完成后延迟一段时间读取状态
        readout_time = self.T + self.delay_steps - 1
        if readout_time >= self.T:
            # 如果延迟超过输入时长，使用最后状态
            final_state = liquid_states[:, -1]
        else:
            final_state = liquid_states[:, readout_time]
            
        # 读出
        output = self.readout(final_state)
        return output

def load_battery_data(file_path, normalize=True, scaler=None):
    """加载电池数据"""
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"数据文件不存在: {file_path}")
    
    df = pd.read_csv(file_path)
    
    # 验证必需列
    required_columns = ['cycle', 'Voltages', 'rate', 'Tem', 'Capacity', 'initial_capacity']
    missing_columns = [col for col in required_columns if col not in df.columns]
    if missing_columns:
        raise ValueError(f"数据文件缺少必需列: {missing_columns}")
    
    # 解析电压序列字符串
    voltages = []
    for i, volt_str in enumerate(df['Voltages']):
        try:
            # 尝试直接解析列表字符串
            volt_list = ast.literal_eval(volt_str)
            if len(volt_list) != 14:
                raise ValueError(f"电压序列长度应为14，但得到{len(volt_list)}")
            voltages.append(volt_list)
        except Exception as e:
            # 如果失败，尝试其他解析方法
            try:
                volt_str_clean = volt_str.strip('[]').split(',')
                volt_list = [float(v.strip()) for v in volt_str_clean]
                if len(volt_list) != 14:
                    raise ValueError(f"电压序列长度应为14，但得到{len(volt_list)}")
                voltages.append(volt_list)
            except Exception as e2:
                print(f"警告: 无法解析第{i+1}行的电压序列: {volt_str}")
                print(f"错误: {e2}")
                continue
    
    # 创建特征矩阵
    X = np.array(voltages)
    
    # 使用容量作为标签
    y = df['Capacity'].values
    
    # 获取初始容量
    initial_capacities = df['initial_capacity'].values
    
    # 数据验证
    if len(X) == 0:
        raise ValueError("没有成功解析任何电压序列数据")
    
    if len(X) != len(y):
        raise ValueError("特征和标签数量不匹配")
    
    if len(X) != len(initial_capacities):
        raise ValueError("特征和初始容量数量不匹配")
    
    # 数据标准化
    if normalize:
        if scaler is None:
            scaler = StandardScaler()
            X = scaler.fit_transform(X)
            return X, y, initial_capacities, scaler
        else:
            X = scaler.transform(X)
            return X, y, initial_capacities, scaler
    
    return X, y, initial_capacities, None

def create_data_loader(X, y, batch_size=32, shuffle=True):
    """创建数据加载器"""
    X_tensor = torch.FloatTensor(X)
    y_tensor = torch.FloatTensor(y)
    
    dataset = torch.utils.data.TensorDataset(X_tensor, y_tensor)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle
    )
    
    return dataloader

def train_model(model, train_loader, epochs=100, lr=0.01):
    """训练模型"""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()
    
    model.train()
    losses = []
    
    for epoch in range(epochs):
        total_loss = 0
        for batch_x, batch_y in train_loader:
            batch_x = batch_x.to(DEVICE)
            batch_y = batch_y.to(DEVICE)
            
            optimizer.zero_grad()
            output = model(batch_x)
            loss = criterion(output.squeeze(), batch_y)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            
        avg_loss = total_loss / len(train_loader)
        losses.append(avg_loss)
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch}, Loss: {avg_loss:.4f}")
        
        # 重置神经元状态
        functional.reset_net(model)
    
    return losses

def evaluate_model(model, test_loader, initial_capacities=None):
    """评估模型
    Args:
        model: 要评估的模型
        test_loader: 测试数据加载器
        initial_capacities: 初始容量数组，如果提供则使用SOH计算误差
    """
    model.eval()
    predictions = []
    targets = []
    
    with torch.no_grad():
        for batch_x, batch_y in test_loader:
            batch_x = batch_x.to(DEVICE)
            output = model(batch_x)
            
            predictions.extend(output.squeeze().cpu().numpy())
            targets.extend(batch_y.cpu().numpy())
    
    predictions = np.array(predictions)
    targets = np.array(targets)
    
    # 计算评估指标
    if initial_capacities is not None:
        # 使用SOH（容量/初始容量）计算误差
        soh_targets = targets / initial_capacities
        soh_predictions = predictions / initial_capacities
        mse = mean_squared_error(soh_targets, soh_predictions)
        mae = mean_absolute_error(soh_targets, soh_predictions)
        rmse = np.sqrt(mse)
    else:
        # 使用容量计算误差（保持向后兼容）
        mse = mean_squared_error(targets, predictions)
        mae = mean_absolute_error(targets, predictions)
        rmse = np.sqrt(mse)
    
    return {
        'mse': mse,
        'mae': mae, 
        'rmse': rmse,
        'predictions': predictions,
        'targets': targets
    }

def save_model_weights(model, file_path):
    """保存模型权重"""
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    torch.save(model.state_dict(), file_path)
    print(f"模型权重已保存到: {file_path}")

def load_model_weights(model, file_path):
    """加载模型权重"""
    if os.path.exists(file_path):
        model.load_state_dict(torch.load(file_path, map_location=DEVICE))
        print(f"模型权重已从 {file_path} 加载")
        return True
    else:
        print(f"权重文件 {file_path} 不存在")
        return False

def main():
    """主程序"""
    parser = argparse.ArgumentParser(description='LSM电池容量预测模型')
    parser.add_argument('--mode', type=str, default='train_test', 
                       choices=['train', 'test', 'train_test'],
                       help='运行模式: train/test/train_test')
    parser.add_argument('--train_data', type=str, 
                       default='data_set/NCA_train_data.csv',
                       help='训练数据文件路径')
    parser.add_argument('--test_data', type=str,
                       default='data_set/NCA_test_data.csv', 
                       help='测试数据文件路径')
    parser.add_argument('--epochs', type=int, default=100,
                       help='训练轮数')
    parser.add_argument('--lr', type=float, default=0.01,
                       help='学习率')
    parser.add_argument('--batch_size', type=int, default=32,
                       help='批次大小')
    parser.add_argument('--weights_file', type=str,
                       default='weight/lsm_battery_weights.pth',
                       help='权重文件路径')
    
    args = parser.parse_args()
    
    # 初始化模型
    model = BatteryLSM(
        input_size=INPUT_SIZE,
        liquid_neurons=LIQUID_NEURONS,
        readout_neurons=READOUT_NEURONS,
        T=TIME_STEPS,
        delay_steps=DELAY_STEPS,
        params=LSM_PARAMS
    ).to(DEVICE)
    
    print(f"使用设备: {DEVICE}")
    print(f"模型参数: {sum(p.numel() for p in model.parameters())}")
    
    if args.mode in ['train', 'train_test']:
        # 训练模式
        print(f"加载训练数据: {args.train_data}")
        X_train, y_train, initial_capacities_train, scaler = load_battery_data(args.train_data)
        train_loader = create_data_loader(X_train, y_train, batch_size=args.batch_size)
        
        print(f"开始训练，轮数: {args.epochs}, 学习率: {args.lr}")
        train_model(model, train_loader, epochs=args.epochs, lr=args.lr)
        
        # 保存权重
        save_model_weights(model, args.weights_file)
    
    if args.mode in ['test', 'train_test']:
        # 测试模式
        if args.mode == 'test':
            # 仅测试模式需要加载权重
            if not load_model_weights(model, args.weights_file):
                print("无法加载权重文件，退出测试模式")
                return
        
        print(f"加载测试数据: {args.test_data}")
        X_test, y_test, initial_capacities_test, _ = load_battery_data(args.test_data)
        test_loader = create_data_loader(X_test, y_test, batch_size=args.batch_size, shuffle=False)
        
        print("开始测试...")
        # 使用SOH计算误差
        results = evaluate_model(model, test_loader, initial_capacities_test)
        
        print("\n=== 测试结果 (基于SOH计算) ===")
        print(f"MSE: {results['mse']:.6f}")
        print(f"MAE: {results['mae']:.6f}")
        print(f"RMSE: {results['rmse']:.6f}")
        
        # 显示前10个样本的预测结果（容量值）
        print("\n前10个样本预测结果 (容量值):")
        print("真实容量 -> 预测容量")
        for i in range(min(10, len(results['targets']))):
            print(f"{results['targets'][i]:.2f} -> {results['predictions'][i]:.2f}")
        
        # 显示前10个样本的SOH值
        print("\n前10个样本SOH值:")
        print("真实SOH -> 预测SOH")
        for i in range(min(10, len(results['targets']))):
            true_soh = results['targets'][i] / initial_capacities_test[i]
            pred_soh = results['predictions'][i] / initial_capacities_test[i]
            print(f"{true_soh:.4f} -> {pred_soh:.4f}")

if __name__ == "__main__":
    main()
